# 0. 引入必要的包
import os
from sklearn.preprocessing import LabelEncoder
import numpy as np
from sklearn.model_selection import train_test_split
from tqdm import tqdm

# TODO
from util import get, preprocess_image,dump,load

# 1. 读取配置文件中的信息
train_dir = get("train") # 获取 训练数据路径
char_styles = get("char_styles") # 获取 字符样式列表，注意: 必须是列标
new_size = get("new_size") # 获取 新图像大小元组, 注意: 必须包含h和w
print(train_dir, char_styles, new_size)
# 2. 生成X,y
print("# 读取训练数据并进行预处理，")
#定义一个函数，根据指定目录、书法体和图像尺寸创建X和y数组
def create_X_y(data_dir, char_styles, new_size):
    # 获取指定目录中的所有图像路径
    image_paths = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if os.path.isfile(os.path.join(data_dir, f))]

    # 创建X和y数组
    X = []
    y = []
    num_images = len(image_paths)

    # 按照书法体分组
    style_counts = {style: 0 for style in char_styles}# 初始化书法体的计数器
    for image_path in image_paths:
        label = image_path.split(os.path.sep)[-1].split("_")[1]   # 从图像路径中获取书法体标签
        style_counts[label] += 1   # 书法体计数器加1

    label_encoder = LabelEncoder()   # 创建标签编码器
    label_encoder.fit(char_styles)   # 标签编码器拟合书法体标签
    # 为每个书法体分别读取图像并添加到X和y数组中
    for style in char_styles:
        style_paths = [image_path for image_path in image_paths if
                       image_path.split(os.path.sep)[-1].split("_")[1] == style]  # 获取所有属于当前书法体的图像路径
        for image_path in tqdm(style_paths, desc=f"处理 {style} 图像", ncols=100):  # 使用tqdm显示进度条
            pixels = preprocess_image(image_path, new_size)  # 读取图像并进行预处理
            X.append(pixels)  # 将预处理后的图像数据添加到X数组中
            y.append(label_encoder.transform([style])[0])  # 将当前书法体标签的编码添加到y数组中

    return np.array(X), np.array(y)  # 将X和y数组转换为numpy数组并返回

############################## 程序逻辑部分 ##############################
# 1. 创建/读取X和y
X, y = create_X_y(train_dir, char_styles, new_size)
print("成功读取X和y")
y=y.astype(np.int64)
# TODO
# 3. 分割测试集和训练集
print("# 将数据按 80% 和 20% 的比例分割")
X_train, X_test, y_train, y_test =train_test_split(X, y, test_size=0.2) #TODO

# 4. 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

# 5. 序列化分割后的训练和测试样本
dump((X_train, X_test, y_train, y_test), "(X_train, X_test, y_train, y_test)", f'{get("Xy_root")}/Xy')

# TODO